{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed data parallel BERT training with TensorFlow2 and SMDataParallel\n",
    "\n",
    "HSMDataParallel is a new capability in Amazon SageMaker to train deep learning models faster and cheaper. SMDataParallel is a distributed data parallel training framework for TensorFlow, PyTorch, and MXNet.\n",
    "\n",
    "This notebook example shows how to use SMDataParallel with TensorFlow(version 2.3.1) on [Amazon SageMaker](https://aws.amazon.com/sagemaker/) to train a BERT model using [Amazon FSx for Lustre file-system](https://aws.amazon.com/fsx/lustre/) as data source.\n",
    "\n",
    "The outline of steps is as follows:\n",
    "\n",
    "1. Stage dataset in [Amazon S3](https://aws.amazon.com/s3/). Original dataset for BERT pretraining consists of text passages from BooksCorpus (800M words) (Zhu et al. 2015) and English Wikipedia (2,500M words). Please follow original guidelines by NVidia to prepare training data in hdf5 format - \n",
    "https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/README.md#getting-the-data\n",
    "2. Create Amazon FSx Lustre file-system and import data into the file-system from S3\n",
    "3. Build Docker training image and push it to [Amazon ECR](https://aws.amazon.com/ecr/)\n",
    "4. Configure data input channels for SageMaker\n",
    "5. Configure hyper-prarameters\n",
    "6. Define training metrics\n",
    "7. Define training job, set distribution strategy to SMDataParallel and start training\n",
    "\n",
    "**NOTE:**  With large traning dataset, we recommend using (Amazon FSx)[https://aws.amazon.com/fsx/] as the input filesystem for the SageMaker training job. FSx file input to SageMaker significantly cuts down training start up time on SageMaker because it avoids downloading the training data each time you start the training job (as done with S3 input for SageMaker training job) and provides good data read throughput.\n",
    "\n",
    "\n",
    "**NOTE:** This example requires SageMaker Python SDK v2.X."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Amazon SageMaker Initialization\n",
    "\n",
    "Initialize the notebook instance. Get the aws region, sagemaker execution role.\n",
    "\n",
    "The IAM role arn used to give training and hosting access to your data. See the [Amazon SageMaker Roles](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) for how to create these. Note, if more than one role is required for notebook instances, training, and/or hosting, please replace the sagemaker.get_execution_role() with the appropriate full IAM role arn string(s). As described above, since we will be using FSx, please make sure to attach `FSx Access` permission to this IAM role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already up-to-date: sagemaker in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (2.19.0)\n",
      "Requirement already satisfied, skipping upgrade: boto3>=1.16.32 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (1.16.33)\n",
      "Requirement already satisfied, skipping upgrade: google-pasta in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (0.2.0)\n",
      "Requirement already satisfied, skipping upgrade: protobuf>=3.1 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (3.8.0)\n",
      "Requirement already satisfied, skipping upgrade: packaging>=20.0 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (20.1)\n",
      "Requirement already satisfied, skipping upgrade: importlib-metadata>=1.4.0 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (1.5.0)\n",
      "Requirement already satisfied, skipping upgrade: protobuf3-to-dict>=0.1.5 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (0.1.5)\n",
      "Requirement already satisfied, skipping upgrade: attrs in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (19.3.0)\n",
      "Requirement already satisfied, skipping upgrade: numpy>=1.9.0 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (1.18.1)\n",
      "Requirement already satisfied, skipping upgrade: smdebug-rulesconfig>=1.0.0 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from sagemaker) (1.0.0)\n",
      "Requirement already satisfied, skipping upgrade: botocore<1.20.0,>=1.19.33 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from boto3>=1.16.32->sagemaker) (1.19.33)\n",
      "Requirement already satisfied, skipping upgrade: s3transfer<0.4.0,>=0.3.0 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from boto3>=1.16.32->sagemaker) (0.3.3)\n",
      "Requirement already satisfied, skipping upgrade: jmespath<1.0.0,>=0.7.1 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from boto3>=1.16.32->sagemaker) (0.10.0)\n",
      "Requirement already satisfied, skipping upgrade: six in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from google-pasta->sagemaker) (1.14.0)\n",
      "Requirement already satisfied, skipping upgrade: setuptools in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from protobuf>=3.1->sagemaker) (45.2.0.post20200210)\n",
      "Requirement already satisfied, skipping upgrade: pyparsing>=2.0.2 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from packaging>=20.0->sagemaker) (2.4.6)\n",
      "Requirement already satisfied, skipping upgrade: zipp>=0.5 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from importlib-metadata>=1.4.0->sagemaker) (2.2.0)\n",
      "Requirement already satisfied, skipping upgrade: urllib3<1.27,>=1.25.4; python_version != \"3.4\" in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from botocore<1.20.0,>=1.19.33->boto3>=1.16.32->sagemaker) (1.25.10)\n",
      "Requirement already satisfied, skipping upgrade: python-dateutil<3.0.0,>=2.1 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from botocore<1.20.0,>=1.19.33->boto3>=1.16.32->sagemaker) (2.8.1)\n",
      "\u001b[33mWARNING: You are using pip version 20.0.2; however, version 20.3.1 is available.\n",
      "You should consider upgrading via the '/home/ec2-user/anaconda3/envs/tensorflow_p36/bin/python3 -m pip install --upgrade pip' command.\u001b[0m\n",
      "SageMaker Execution Role:arn:aws:iam::835319576252:role/service-role/AmazonSageMaker-ExecutionRole-20191006T135881\n",
      "AWS account:835319576252\n",
      "AWS region:us-east-1\n",
      "CPU times: user 122 ms, sys: 12.2 ms, total: 134 ms\n",
      "Wall time: 1.87 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "! python3 -m pip install --upgrade sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SageMaker Execution Role:arn:aws:iam::835319576252:role/service-role/AmazonSageMaker-ExecutionRole-20191006T135881\n",
      "AWS account:835319576252\n",
      "AWS region:us-east-1\n"
     ]
    }
   ],
   "source": [
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.estimator import Estimator\n",
    "import boto3\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "\n",
    "role = get_execution_role() # provide a pre-existing role ARN as an alternative to creating a new role\n",
    "print(f'SageMaker Execution Role:{role}')\n",
    "\n",
    "client = boto3.client('sts')\n",
    "account = client.get_caller_identity()['Account']\n",
    "print(f'AWS account:{account}')\n",
    "\n",
    "session = boto3.session.Session()\n",
    "region = session.region_name\n",
    "print(f'AWS region:{region}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare SageMaker Training Images\n",
    "\n",
    "1. SageMaker by default use the latest [Amazon Deep Learning Container Images (DLC)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md) TensorFlow training image. In this step, we use it as a base image and install additional dependencies required for training BERT model.\n",
    "2. In the Github repository https://github.com/HerringForks/DeepLearningExamples.git we have made TensorFlow2-SMDataParallel BERT training script available for your use. This repository will be cloned in the training image for running the model training.\n",
    "\n",
    "### Build and Push Docker Image to ECR\n",
    "\n",
    "Run the below command build the docker image and push it to ECR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = \"tf2-smdataparallel-bert-sagemaker\"  # Example: tf2-smdataparallel-bert-sagemaker\n",
    "tag = \"latest\"   # Example: latest "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mARG\u001b[39;49;00m region\r\n",
      "\r\n",
      "\u001b[34mFROM\u001b[39;49;00m \u001b[33m763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:2.3.1-gpu-py37-cu110-ubuntu18.04\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[34mRUN\u001b[39;49;00m \tpip --no-cache-dir --no-cache install \u001b[33m\\\u001b[39;49;00m\r\n",
      "        scikit-learn==\u001b[34m0\u001b[39;49;00m.23.1 \u001b[33m\\\u001b[39;49;00m\r\n",
      "        \u001b[31mwandb\u001b[39;49;00m==\u001b[34m0\u001b[39;49;00m.9.1 \u001b[33m\\\u001b[39;49;00m\r\n",
      "        tensorflow-addons \u001b[33m\\\u001b[39;49;00m\r\n",
      "        \u001b[31mcolorama\u001b[39;49;00m==\u001b[34m0\u001b[39;49;00m.4.3 \u001b[33m\\\u001b[39;49;00m\r\n",
      "        pandas \u001b[33m\\\u001b[39;49;00m\r\n",
      "        apache_beam \u001b[33m\\\u001b[39;49;00m\r\n",
      "        \u001b[31mpyarrow\u001b[39;49;00m==\u001b[34m0\u001b[39;49;00m.16 \u001b[33m\\\u001b[39;49;00m\r\n",
      "        git+https://github.com/HerringForks/transformers.git@master \u001b[33m\\\u001b[39;49;00m\r\n",
      "        git+https://github.com/huggingface/nlp.git@703b761\r\n",
      "        \r\n"
     ]
    }
   ],
   "source": [
    "!pygmentize ./Dockerfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37m#!/usr/bin/env bash\u001b[39;49;00m\r\n",
      "\u001b[37m# This script shows how to build the Docker image and push it to ECR to be ready for use\u001b[39;49;00m\r\n",
      "\u001b[37m# by SageMaker.\u001b[39;49;00m\r\n",
      "\u001b[37m# The argument to this script is the image name. This will be used as the image on the local\u001b[39;49;00m\r\n",
      "\u001b[37m# machine and combined with the account and region to form the repository name for ECR.\u001b[39;49;00m\r\n",
      "\u001b[37m# set region\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[31mDIR\u001b[39;49;00m=\u001b[33m\"\u001b[39;49;00m\u001b[34m$(\u001b[39;49;00m \u001b[36mcd\u001b[39;49;00m \u001b[33m\"\u001b[39;49;00m\u001b[34m$(\u001b[39;49;00m dirname \u001b[33m\"\u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31mBASH_SOURCE\u001b[39;49;00m[0]\u001b[33m}\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m \u001b[34m)\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m && \u001b[36mpwd\u001b[39;49;00m \u001b[34m)\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[34mif\u001b[39;49;00m [ \u001b[33m\"\u001b[39;49;00m\u001b[31m$#\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m -eq \u001b[34m3\u001b[39;49;00m ]; \u001b[34mthen\u001b[39;49;00m\r\n",
      "    \u001b[31mregion\u001b[39;49;00m=\u001b[31m$1\u001b[39;49;00m\r\n",
      "    \u001b[31mimage\u001b[39;49;00m=\u001b[31m$2\u001b[39;49;00m\r\n",
      "    \u001b[31mtag\u001b[39;49;00m=\u001b[31m$3\u001b[39;49;00m\r\n",
      "\u001b[34melse\u001b[39;49;00m\r\n",
      "    \u001b[36mecho\u001b[39;49;00m \u001b[33m\"\u001b[39;49;00m\u001b[33musage: \u001b[39;49;00m\u001b[31m$0\u001b[39;49;00m\u001b[33m <aws-region> \u001b[39;49;00m\u001b[31m$1\u001b[39;49;00m\u001b[33m <image-repo> \u001b[39;49;00m\u001b[31m$2\u001b[39;49;00m\u001b[33m <image-tag>\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m\r\n",
      "    \u001b[36mexit\u001b[39;49;00m \u001b[34m1\u001b[39;49;00m\r\n",
      "\u001b[34mfi\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[37m# Get the account number associated with the current IAM credentials\u001b[39;49;00m\r\n",
      "\u001b[31maccount\u001b[39;49;00m=\u001b[34m$(\u001b[39;49;00maws sts get-caller-identity --query Account --output text\u001b[34m)\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[34mif\u001b[39;49;00m [ \u001b[31m$?\u001b[39;49;00m -ne \u001b[34m0\u001b[39;49;00m ]\r\n",
      "\u001b[34mthen\u001b[39;49;00m\r\n",
      "    \u001b[36mexit\u001b[39;49;00m \u001b[34m255\u001b[39;49;00m\r\n",
      "\u001b[34mfi\u001b[39;49;00m\r\n",
      "\r\n",
      "\r\n",
      "\u001b[31mfullname\u001b[39;49;00m=\u001b[33m\"\u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31maccount\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\u001b[33m.dkr.ecr.\u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31mregion\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\u001b[33m.amazonaws.com/\u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31mimage\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\u001b[33m:\u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31mtag\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[37m# If the repository doesn't exist in ECR, create it.\u001b[39;49;00m\r\n",
      "aws ecr describe-repositories --region \u001b[33m${\u001b[39;49;00m\u001b[31mregion\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m --repository-names \u001b[33m\"\u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31mimage\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m > /dev/null \u001b[34m2\u001b[39;49;00m>&\u001b[34m1\u001b[39;49;00m\r\n",
      "\u001b[34mif\u001b[39;49;00m [ \u001b[31m$?\u001b[39;49;00m -ne \u001b[34m0\u001b[39;49;00m ]; \u001b[34mthen\u001b[39;49;00m\r\n",
      "    \u001b[36mecho\u001b[39;49;00m \u001b[33m\"\u001b[39;49;00m\u001b[33mcreating ECR repository : \u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31mfullname\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\u001b[33m \u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m\r\n",
      "    aws ecr create-repository --region \u001b[33m${\u001b[39;49;00m\u001b[31mregion\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m --repository-name \u001b[33m\"\u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31mimage\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m > /dev/null\r\n",
      "\u001b[34mfi\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[34m$(\u001b[39;49;00maws ecr get-login --no-include-email --region us-west-2  --registry-ids \u001b[34m763104351884\u001b[39;49;00m\u001b[34m)\u001b[39;49;00m\r\n",
      "docker build \u001b[33m${\u001b[39;49;00m\u001b[31mDIR\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m/ -t \u001b[33m${\u001b[39;49;00m\u001b[31mimage\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m -f \u001b[33m${\u001b[39;49;00m\u001b[31mDIR\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m/Dockerfile  --build-arg \u001b[31mregion\u001b[39;49;00m=\u001b[33m${\u001b[39;49;00m\u001b[31mregion\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\r\n",
      "docker tag \u001b[33m${\u001b[39;49;00m\u001b[31mimage\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m \u001b[33m${\u001b[39;49;00m\u001b[31mfullname\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[37m# Get the login command from ECR and execute it directly\u001b[39;49;00m\r\n",
      "\u001b[34m$(\u001b[39;49;00maws ecr get-login --region \u001b[33m${\u001b[39;49;00m\u001b[31mregion\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m --no-include-email\u001b[34m)\u001b[39;49;00m\r\n",
      "docker push \u001b[33m${\u001b[39;49;00m\u001b[31mfullname\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\r\n",
      "\u001b[34mif\u001b[39;49;00m [ \u001b[31m$?\u001b[39;49;00m -eq \u001b[34m0\u001b[39;49;00m ]; \u001b[34mthen\u001b[39;49;00m\r\n",
      "\t\u001b[36mecho\u001b[39;49;00m \u001b[33m\"\u001b[39;49;00m\u001b[33mAmazon ECR URI: \u001b[39;49;00m\u001b[33m${\u001b[39;49;00m\u001b[31mfullname\u001b[39;49;00m\u001b[33m}\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m\r\n",
      "\u001b[34melse\u001b[39;49;00m\r\n",
      "\t\u001b[36mecho\u001b[39;49;00m \u001b[33m\"Error: Image build and push failed\"\u001b[39;49;00m\r\n",
      "\t\u001b[36mexit\u001b[39;49;00m \u001b[34m1\u001b[39;49;00m\r\n",
      "\u001b[34mfi\u001b[39;49;00m\r\n"
     ]
    }
   ],
   "source": [
    "!pygmentize ./build_and_push.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING! Using --password via the CLI is insecure. Use --password-stdin.\n",
      "WARNING! Your password will be stored unencrypted in /home/ec2-user/.docker/config.json.\n",
      "Configure a credential helper to remove this warning. See\n",
      "https://docs.docker.com/engine/reference/commandline/login/#credentials-store\n",
      "\n",
      "Login Succeeded\n",
      "Sending build context to Docker daemon  12.35MB\n",
      "Step 1/3 : ARG region\n",
      "Step 2/3 : FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:2.3.1-gpu-py37-cu110-ubuntu18.04\n",
      " ---> 73f448953d3a\n",
      "Step 3/3 : RUN \tpip --no-cache-dir --no-cache install         scikit-learn==0.23.1         wandb==0.9.1         tensorflow-addons         colorama==0.4.3         pandas         apache_beam         pyarrow==0.16         git+https://github.com/HerringForks/transformers.git@master         git+https://github.com/huggingface/nlp.git@703b761\n",
      " ---> Using cache\n",
      " ---> 24901ecc9de0\n",
      "Successfully built 24901ecc9de0\n",
      "Successfully tagged tf2-smdataparallel-bert-sagemaker:latest\n",
      "WARNING! Using --password via the CLI is insecure. Use --password-stdin.\n",
      "WARNING! Your password will be stored unencrypted in /home/ec2-user/.docker/config.json.\n",
      "Configure a credential helper to remove this warning. See\n",
      "https://docs.docker.com/engine/reference/commandline/login/#credentials-store\n",
      "\n",
      "Login Succeeded\n",
      "The push refers to repository [835319576252.dkr.ecr.us-east-1.amazonaws.com/tf2-smdataparallel-bert-sagemaker]\n",
      "\n",
      "\u001b[1Bc96185e7: Preparing \n",
      "\u001b[1Bda29752b: Preparing \n",
      "\u001b[1Baf7081a0: Preparing \n",
      "\u001b[1B8f790f4a: Preparing \n",
      "\u001b[1Ba172edf1: Preparing \n",
      "\u001b[1B75917cc2: Preparing \n",
      "\u001b[1B445f2bde: Preparing \n",
      "\u001b[1B7f80b0c1: Preparing \n",
      "\u001b[1B90dd5fd5: Preparing \n",
      "\u001b[1B98b178d8: Preparing \n",
      "\u001b[1B7efd8228: Preparing \n",
      "\u001b[1Bdb20b91f: Preparing \n",
      "\u001b[1Bb9afb958: Preparing \n",
      "\u001b[1B7406ed5a: Preparing \n",
      "\u001b[1Bde285e41: Preparing \n",
      "\u001b[1B581d23ab: Preparing \n",
      "\u001b[1B8072e44c: Preparing \n",
      "\u001b[1Bca2d4d4f: Preparing \n",
      "\u001b[1Bf8c1572c: Preparing \n",
      "\u001b[14B45f2bde: Waiting g \n",
      "\u001b[13B0dd5fd5: Waiting g \n",
      "\u001b[1Bc9c88d14: Preparing \n",
      "\u001b[14B8b178d8: Waiting g \n",
      "\u001b[1B38a9aef5: Preparing \n",
      "\u001b[8Bca2d4d4f: Waiting g \n",
      "\u001b[11B81d23ab: Waiting g \n",
      "\u001b[11B072e44c: Waiting g \n",
      "\u001b[14Be285e41: Waiting g \n",
      "\u001b[16B406ed5a: Waiting g \n",
      "\u001b[10Bef50ffc: Waiting g \n",
      "\u001b[10B9c88d14: Waiting g \n",
      "\u001b[9B38a9aef5: Waiting g \n",
      "\u001b[1B34ed948a: Preparing \n",
      "\u001b[6B1cf2580d: Waiting g \n",
      "\u001b[6B604caecc: Waiting g \n",
      "\u001b[10B1d66695: Waiting g \n",
      "\u001b[1B93ac9526: Preparing \n",
      "\u001b[11B44c8522: Waiting g \n",
      "\u001b[34B5917cc2: Waiting g \n",
      "\u001b[1B824934b3: Preparing \n",
      "\u001b[1B9d46081a: Preparing \n",
      "\u001b[12B7c45380: Waiting g \n",
      "\u001b[12Be7b68e0: Waiting g \n",
      "\u001b[1B67539a0c: Preparing \n",
      "\u001b[13B4ed948a: Waiting g \n",
      "\u001b[1Ba984e1d1: Preparing \n",
      "\u001b[14Ba36ae91: Waiting g \n",
      "\u001b[1Bf94c75c1: Preparing \n",
      "\u001b[1Bb14f6ff1: Preparing \n",
      "\u001b[1B5f19e930: Preparing \n",
      "\u001b[1B26f8b991: Preparing \n",
      "\u001b[18Be003d68: Waiting g \n",
      "\u001b[1B4dce1444: Preparing \n",
      "\u001b[19B85beea2: Waiting g \n",
      "\u001b[1Be116c0c0: Preparing \n",
      "\u001b[1B4df0ad6c: Preparing \n",
      "\u001b[15Be1c7223: Waiting g \n",
      "\u001b[1B02706667: Layer already exists \u001b[53A\u001b[2K\u001b[48A\u001b[2K\u001b[47A\u001b[2K\u001b[39A\u001b[2K\u001b[35A\u001b[2K\u001b[31A\u001b[2K\u001b[27A\u001b[2K\u001b[22A\u001b[2K\u001b[16A\u001b[2K\u001b[12A\u001b[2K\u001b[9A\u001b[2K\u001b[4A\u001b[2Klatest: digest: sha256:a0e36b294f1909845a48d8e14725b922a137669bcb13c19e2f1029381f3c216d size: 12499\n",
      "Amazon ECR URI: 835319576252.dkr.ecr.us-east-1.amazonaws.com/tf2-smdataparallel-bert-sagemaker:latest\n",
      "CPU times: user 88.4 ms, sys: 18.1 ms, total: 106 ms\n",
      "Wall time: 8.34 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "! chmod +x build_and_push.sh; bash build_and_push.sh {region} {image} {tag}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing FSx Input for SageMaker\n",
    "\n",
    "1. Download and prepare your training dataset on S3.\n",
    "2. Follow the steps listed here to create a FSx linked with your S3 bucket with training data - https://docs.aws.amazon.com/fsx/latest/LustreGuide/create-fs-linked-data-repo.html. Make sure to add an endpoint to your VPC allowing S3 access.\n",
    "3. Follow the steps listed here to configure your SageMaker training job to use FSx https://aws.amazon.com/blogs/machine-learning/speed-up-training-on-amazon-sagemaker-using-amazon-efs-or-amazon-fsx-for-lustre-file-systems/\n",
    "\n",
    "### Important Caveats\n",
    "\n",
    "1. You need use the same `subnet` and `vpc` and `security group` used with FSx when launching the SageMaker notebook instance. The same configurations will be used by your SageMaker training job.\n",
    "2. Make sure you set appropriate inbound/output rules in the `security group`. Specically, opening up these ports is necessary for SageMaker to access the FSx filesystem in the training job. https://docs.aws.amazon.com/fsx/latest/LustreGuide/limit-access-security-groups.html\n",
    "3. Make sure `SageMaker IAM Role` used to launch this SageMaker training job has access to `AmazonFSx`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SageMaker TensorFlow Estimator function options\n",
    "\n",
    "In the following code block, you can update the estimator function to use a different instance type, instance count, and distrubtion strategy. You're also passing in the training script you reviewed in the previous cell.\n",
    "\n",
    "**Instance types**\n",
    "\n",
    "SMDataParallel supports model training on SageMaker with the following instance types only:\n",
    "1. ml.p3.16xlarge\n",
    "1. ml.p3dn.24xlarge [Recommended]\n",
    "1. ml.p4d.24xlarge [Recommended]\n",
    "\n",
    "**Instance count**\n",
    "\n",
    "To get the best performance and the most out of SMDataParallel, you should use at least 2 instances, but you can also use 1 for testing this example.\n",
    "\n",
    "**Distribution strategy**\n",
    "\n",
    "Note that to use DDP mode, you update the the `distribution` strategy, and set it to use `smdistributed dataparallel`.\n",
    "\n",
    "### Training script\n",
    "\n",
    "In the Github repository https://github.com/HerringForks/deep-learning-models.git we have made reference TensorFlow-SMDataParallel BERT training script available for your use. Clone the repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'deep-learning-models'...\n",
      "remote: Enumerating objects: 42, done.\u001b[K\n",
      "remote: Counting objects: 100% (42/42), done.\u001b[K\n",
      "remote: Compressing objects: 100% (34/34), done.\u001b[K\n",
      "remote: Total 1764 (delta 20), reused 18 (delta 8), pack-reused 1722\u001b[K\n",
      "Receiving objects: 100% (1764/1764), 4.76 MiB | 59.42 MiB/s, done.\n",
      "Resolving deltas: 100% (834/834), done.\n"
     ]
    }
   ],
   "source": [
    "# Clone herring forks repository for reference implementation BERT with TensorFlow2-SMDataParallel\n",
    "!rm -rf deep-learning-models\n",
    "!git clone --recursive https://github.com/HerringForks/deep-learning-models.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "sm = boto3.client('sagemaker')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dsoaws\n"
     ]
    }
   ],
   "source": [
    "notebook_instance_name = sm.list_notebook_instances()['NotebookInstances'][3]['NotebookInstanceName']\n",
    "print(notebook_instance_name)\n",
    "\n",
    "if notebook_instance_name != 'dsoaws':\n",
    "    print('****** ERROR:  MUST FIND THE CORRECT NOTEBOOK ******')\n",
    "    exit()    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'NotebookInstanceArn': 'arn:aws:sagemaker:us-east-1:835319576252:notebook-instance/dsoaws',\n",
       " 'NotebookInstanceName': 'dsoaws',\n",
       " 'NotebookInstanceStatus': 'InService',\n",
       " 'Url': 'dsoaws.notebook.us-east-1.sagemaker.aws',\n",
       " 'InstanceType': 'ml.c5.2xlarge',\n",
       " 'SubnetId': 'subnet-0b8d836c',\n",
       " 'SecurityGroups': ['sg-5383e807'],\n",
       " 'RoleArn': 'arn:aws:iam::835319576252:role/service-role/AmazonSageMaker-ExecutionRole-20191006T135881',\n",
       " 'NetworkInterfaceId': 'eni-0d92cf87d27f516fc',\n",
       " 'LastModifiedTime': datetime.datetime(2020, 11, 28, 4, 35, 41, 853000, tzinfo=tzlocal()),\n",
       " 'CreationTime': datetime.datetime(2020, 2, 24, 23, 6, 31, 851000, tzinfo=tzlocal()),\n",
       " 'DirectInternetAccess': 'Enabled',\n",
       " 'VolumeSizeInGB': 2000,\n",
       " 'RootAccess': 'Enabled',\n",
       " 'ResponseMetadata': {'RequestId': 'eb8ddbd7-d4c7-403a-9d57-dce93069c1d9',\n",
       "  'HTTPStatusCode': 200,\n",
       "  'HTTPHeaders': {'x-amzn-requestid': 'eb8ddbd7-d4c7-403a-9d57-dce93069c1d9',\n",
       "   'content-type': 'application/x-amz-json-1.1',\n",
       "   'content-length': '594',\n",
       "   'date': 'Thu, 10 Dec 2020 20:04:58 GMT'},\n",
       "  'RetryAttempts': 0}}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "notebook_instance = sm.describe_notebook_instance(NotebookInstanceName=notebook_instance_name)\n",
    "notebook_instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sg-5383e807\n"
     ]
    }
   ],
   "source": [
    "security_group_id = notebook_instance['SecurityGroups'][0]\n",
    "print(security_group_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subnet-0b8d836c\n"
     ]
    }
   ],
   "source": [
    "subnet_id = notebook_instance['SubnetId']\n",
    "print(subnet_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tensorflow import TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "835319576252\n",
      "us-east-1\n",
      "tf2-smdataparallel-bert-sagemaker\n",
      "latest\n"
     ]
    }
   ],
   "source": [
    "print(account)\n",
    "print(region)\n",
    "print(image)\n",
    "print(tag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_type = \"ml.p3dn.24xlarge\" # Other supported instance type: ml.p3.16xlarge, ml.p4d.24xlarge\n",
    "instance_count = 2 # You can use 2, 4, 8 etc.\n",
    "docker_image = f\"{account}.dkr.ecr.{region}.amazonaws.com/{image}:{tag}\" # YOUR_ECR_IMAGE_BUILT_WITH_ABOVE_DOCKER_FILE\n",
    "username = 'AWS'\n",
    "subnets = [subnet_id] # Should be same as Subnet used for FSx. Example: subnet-0f9XXXX\n",
    "security_group_ids = [security_group_id] # Should be same as Security group used for FSx. sg-03ZZZZZZ\n",
    "job_name = 'smdataparallel-bert-tf2-fsx-2p3dn' # This job name is used as prefix to the sagemaker training job. Makes it easy for your look for your training job in SageMaker Training job console.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO:  Copy data to FSx/S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting datasets\n",
      "  Downloading datasets-1.1.3-py3-none-any.whl (153 kB)\n",
      "\u001b[K     |████████████████████████████████| 153 kB 10.7 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pyarrow>=0.17.1\n",
      "  Downloading pyarrow-2.0.0-cp36-cp36m-manylinux2014_x86_64.whl (17.7 MB)\n",
      "\u001b[K     |████████████████████████████████| 17.7 MB 29.0 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: tqdm<4.50.0,>=4.27 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from datasets) (4.42.1)\n",
      "Collecting dill\n",
      "  Downloading dill-0.3.3-py2.py3-none-any.whl (81 kB)\n",
      "\u001b[K     |████████████████████████████████| 81 kB 22.9 MB/s  eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: requests>=2.19.0 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from datasets) (2.22.0)\n",
      "Requirement already satisfied: pandas in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from datasets) (1.0.1)\n",
      "Collecting dataclasses; python_version < \"3.7\"\n",
      "  Downloading dataclasses-0.8-py3-none-any.whl (19 kB)\n",
      "Collecting xxhash\n",
      "  Downloading xxhash-2.0.0-cp36-cp36m-manylinux2010_x86_64.whl (242 kB)\n",
      "\u001b[K     |████████████████████████████████| 242 kB 112.7 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting multiprocess\n",
      "  Downloading multiprocess-0.70.11.1-py36-none-any.whl (101 kB)\n",
      "\u001b[K     |████████████████████████████████| 101 kB 22.0 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: numpy>=1.17 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from datasets) (1.18.1)\n",
      "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from requests>=2.19.0->datasets) (3.0.4)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from requests>=2.19.0->datasets) (1.25.10)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from requests>=2.19.0->datasets) (2020.6.20)\n",
      "Requirement already satisfied: idna<2.9,>=2.5 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from requests>=2.19.0->datasets) (2.8)\n",
      "Requirement already satisfied: pytz>=2017.2 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from pandas->datasets) (2019.3)\n",
      "Requirement already satisfied: python-dateutil>=2.6.1 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from pandas->datasets) (2.8.1)\n",
      "Requirement already satisfied: six>=1.5 in /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages (from python-dateutil>=2.6.1->pandas->datasets) (1.14.0)\n",
      "Installing collected packages: pyarrow, dill, dataclasses, xxhash, multiprocess, datasets\n",
      "Successfully installed dataclasses-0.8 datasets-1.1.3 dill-0.3.3 multiprocess-0.70.11.1 pyarrow-2.0.0 xxhash-2.0.0\n",
      "\u001b[33mWARNING: You are using pip version 20.0.2; however, version 20.3.1 is available.\n",
      "You should consider upgrading via the '/home/ec2-user/anaconda3/envs/tensorflow_p36/bin/python -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/ec2-user/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow_core/__init__.py:1473: The name tf.estimator.inputs is deprecated. Please use tf.compat.v1.estimator.inputs instead.\n",
      "\n",
      "['aeslc', 'afrikaans_ner_corpus', 'ag_news', 'ai2_arc', 'ajgt_twitter_ar', 'allegro_reviews', 'allocine', 'amazon_reviews_multi', 'amazon_us_reviews', 'ambig_qa', 'amttl', 'anli', 'aqua_rat', 'arcd', 'arsentd_lev', 'art', 'arxiv_dataset', 'aslg_pc12', 'asnq', 'asset', 'autshumato', 'bible_para', 'big_patent', 'billsum', 'biomrc', 'blended_skill_talk', 'blimp', 'blog_authorship_corpus', 'bookcorpus', 'bookcorpusopen', 'boolq', 'break_data', 'bsd_ja_en', 'c3', 'c4', 'cail2018', 'capes', 'cawac', 'cdsc', 'cdt', 'cfq', 'chr_en', 'circa', 'civil_comments', 'clinc_oos', 'clue', 'cmrc2018', 'cnn_dailymail', 'coached_conv_pref', 'coarse_discourse', 'codah', 'code_search_net', 'com_qa', 'common_gen', 'commonsense_qa', 'compguesswhat', 'conceptnet5', 'conll2000', 'conll2002', 'conll2003', 'conv_ai', 'coqa', 'cornell_movie_dialog', 'cos_e', 'cosmos_qa', 'covid_qa_castorini', 'covid_qa_deepset', 'craigslist_bargains', 'crd3', 'crime_and_punish', 'crows_pairs', 'cs_restaurants', 'csv', 'curiosity_dialogs', 'daily_dialog', 'dane', 'danish_political_comments', 'dart', 'dbpedia_14', 'dbrd', 'deal_or_no_dialog', 'definite_pronoun_resolution', 'dengue_filipino', 'dialog_re', 'disaster_response_messages', 'discofuse', 'docred', 'doqa', 'dream', 'drop', 'dyk', 'e2e_nlg', 'e2e_nlg_cleaned', 'ecb', 'eitb_parcc', 'eli5', 'emea', 'emo', 'emotion', 'empathetic_dialogues', 'enriched_web_nlg', 'eraser_multi_rc', 'esnli', 'eth_py150_open', 'euronews', 'event2Mind', 'evidence_infer_treatment', 'exams', 'fake_news_filipino', 'farsi_news', 'fever', 'finer', 'flores', 'flue', 'fquad', 'gap', 'generated_reviews_enth', 'german_legal_entity_recognition', 'germaner', 'germeval_14', 'gigaword', 'glucose', 'glue', 'go_emotions', 'google_wellformed_query', 'grail_qa', 'great_code', 'guardian_authorship', 'gutenberg_time', 'hans', 'hansards', 'hard', 'hate_speech_filipino', 'hausa_voa_ner', 'hausa_voa_topics', 'health_fact', 'hebrew_sentiment', 'hellaswag', 'hkcancor', 'hotpot_qa', 'hybrid_qa', 'hyperpartisan_news_detection', 'id_nergrit_ner', 'imdb', 'imdb_urdu_reviews', 'indic_glue', 'inquisitive_qg', 'isixhosa_ner_corpus', 'isizulu_ner_corpus', 'iwslt2017', 'jeopardy', 'jfleg', 'jnlpba', 'json', 'kannada_news', 'kelm', 'kilt_tasks', 'kilt_wikipedia', 'kor_hate', 'kor_ner', 'kor_nli', 'kor_nlu', 'kor_qpair', 'labr', 'lambada', 'large_spanish_corpus', 'lc_quad', 'lener_br', 'librispeech_lm', 'limit', 'lince', 'linnaeus', 'lm1b', 'lst20', 'math_dataset', 'math_qa', 'matinf', 'mc_taco', 'med_hop', 'medal', 'medical_questions_pairs', 'metooma', 'metrec', 'mkb', 'mkqa', 'mlqa', 'mlsum', 'mocha', 'movie_rationales', 'mrqa', 'ms_marco', 'ms_terms', 'msr_genomics_kbcomp', 'msr_text_compression', 'msr_zhen_translation_parity', 'msra_ner', 'multi_news', 'multi_nli', 'multi_nli_mismatch', 'multi_para_crawl', 'multi_woz_v22', 'multi_x_science_sum', 'mwsc', 'myanmar_news', 'natural_questions', 'ncbi_disease', 'nchlt', 'ncslgr', 'neural_code_search', 'news_commentary', 'newsgroup', 'newsph', 'newsph_nli', 'newsroom', 'nkjp-ner', 'nli_tr', 'norwegian_ner', 'nsmc', 'numer_sense', 'numeric_fused_head', 'oclar', 'offenseval2020_tr', 'onestop_english', 'open_subtitles', 'openbookqa', 'openwebtext', 'opinosis', 'opus100', 'opus_dogc', 'opus_elhuyar', 'opus_euconst', 'opus_finlex', 'opus_fiskmo', 'opus_infopankki', 'opus_memat', 'opus_montenegrinsubs', 'opus_openoffice', 'opus_sardware', 'opus_xhosanavy', 'orange_sum', 'pandas', 'para_crawl', 'paws', 'paws-x', 'pec', 'peoples_daily_ner', 'pg19', 'php', 'piaf', 'pib', 'piqa', 'poem_sentiment', 'polemo2', 'polyglot_ner', 'prachathai67k', 'pragmeval', 'proto_qa', 'psc', 'pubmed_qa', 'py_ast', 'qa4mre', 'qa_zre', 'qangaroo', 'qanta', 'qasc', 'qed', 'qed_amara', 'quac', 'quail', 'quarel', 'quartz', 'quora', 'quoref', 'race', 're_dial', 'recipe_nlg', 'reclor', 'reddit', 'reddit_tifu', 'reuters21578', 'roman_urdu', 'ropes', 'rotten_tomatoes', 'sanskrit_classic', 'scan', 'scb_mt_enth_2020', 'schema_guided_dstc8', 'scicite', 'scielo', 'scientific_papers', 'scifact', 'sciq', 'scitail', 'scitldr', 'search_qa', 'sem_eval_2010_task_8', 'sem_eval_2014_task_1', 'sent_comp', 'sentiment140', 'sepedi_ner', 'sesotho_ner_corpus', 'setimes', 'setswana_ner_corpus', 'sharc', 'sharc_modified', 'simple_questions_v2', 'siswati_ner_corpus', 'smartdata', 'sms_spam', 'snli', 'social_bias_frames', 'social_i_qa', 'sogou_news', 'species_800', 'spider', 'squad', 'squad_es', 'squad_it', 'squad_kor_v1', 'squad_v1_pt', 'squad_v2', 'squadshifts', 'stsb_mt_sv', 'style_change_detection', 'super_glue', 'swag', 'swahili_news', 'swedish_ner_corpus', 'swedish_reviews', 'tab_fact', 'tamilmixsentiment', 'tanzil', 'tashkeela', 'taskmaster1', 'taskmaster2', 'taskmaster3', 'tatoeba', 'ted_hrlr', 'ted_multi', 'telugu_books', 'telugu_news', 'tep_en_fa_para', 'text', 'thainer', 'thaiqa_squad', 'thaisum', 'tilde_model', 'tiny_shakespeare', 'tlc', 'totto', 'trec', 'trivia_qa', 'tsac', 'tunizi', 'tuple_ie', 'turkish_ner', 'turku_ner_corpus', 'tweet_qa', 'tweets_ar_en_parallel', 'tydiqa', 'ubuntu_dialogs_corpus', 'udhr', 'um005', 'un_multi', 'un_pc', 'universal_dependencies', 'urdu_fake_news', 'urdu_sentiment_corpus', 'web_nlg', 'web_of_science', 'web_questions', 'weibo_ner', 'wiki40b', 'wiki_auto', 'wiki_bio', 'wiki_dpr', 'wiki_hop', 'wiki_qa', 'wiki_qa_ar', 'wiki_snippets', 'wiki_split', 'wikiann', 'wikicorpus', 'wikihow', 'wikipedia', 'wikisql', 'wikitext', 'wikitext_tl39', 'winograd_wsc', 'winogrande', 'wiqa', 'wisesight1000', 'wisesight_sentiment', 'wmt14', 'wmt15', 'wmt16', 'wmt17', 'wmt18', 'wmt19', 'wmt_t2t', 'wnut_17', 'wongnai_reviews', 'woz_dialogue', 'x_stance', 'xcopa', 'xglue', 'xnli', 'xor_tydi_qa', 'xquad', 'xquad_r', 'xsum', 'xtreme', 'yahoo_answers_qa', 'yahoo_answers_topics', 'yelp_polarity', 'yelp_review_full', 'yoruba_bbc_topics', 'yoruba_gv_ner', 'zest', 'Fraser/news-category-dataset', 'Fraser/python-lines', 'cdminix/mgb1', 'german-nlp-group/german_common_crawl', 'joelito/ler', 'joelito/sem_eval_2010_task_8', 'k-halid/ar', 'lhoestq/squad', 'mulcyber/europarl-mono', 'piEsposito/br-quad-2.0', 'piEsposito/br_quad_20', 'piEsposito/squad_20_ptbr', 'sshleifer/pseudo_bart_xsum']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3d0562a5ac464c43aaff477156767ce6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=4417.0, style=ProgressStyle(description…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8f9e5b67c2074b73864356bbadae66f2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=6866.0, style=ProgressStyle(description…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "ename": "OSError",
     "evalue": "Not enough disk space. Needed: 34.06 GiB (download: 16.99 GiB, generated: 17.07 GiB, post-processed: Unknown size)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mOSError\u001b[0m                                   Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-33-b255159dbed7>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0;31m# To load a dataset\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m \u001b[0mwiki\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_dataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"wikipedia\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"20200501.en\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msplit\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'train'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/datasets/load.py\u001b[0m in \u001b[0;36mload_dataset\u001b[0;34m(path, name, data_dir, data_files, split, cache_dir, features, download_config, download_mode, ignore_verifications, save_infos, script_version, **config_kwargs)\u001b[0m\n\u001b[1;32m    608\u001b[0m         \u001b[0mdownload_config\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdownload_config\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    609\u001b[0m         \u001b[0mdownload_mode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdownload_mode\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 610\u001b[0;31m         \u001b[0mignore_verifications\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mignore_verifications\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    611\u001b[0m     )\n\u001b[1;32m    612\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/datasets/builder.py\u001b[0m in \u001b[0;36mdownload_and_prepare\u001b[0;34m(self, download_config, download_mode, ignore_verifications, try_from_hf_gcs, dl_manager, **download_and_prepare_kwargs)\u001b[0m\n\u001b[1;32m    458\u001b[0m                             \u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdownload_size\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    459\u001b[0m                             \u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset_size\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 460\u001b[0;31m                             \u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpost_processing_size\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    461\u001b[0m                         )\n\u001b[1;32m    462\u001b[0m                     )\n",
      "\u001b[0;31mOSError\u001b[0m: Not enough disk space. Needed: 34.06 GiB (download: 16.99 GiB, generated: 17.07 GiB, post-processed: Unknown size)"
     ]
    }
   ],
   "source": [
    "# For loading datasets\n",
    "from datasets import list_datasets, load_dataset\n",
    "\n",
    "# To see all available dataset names\n",
    "print(list_datasets()) \n",
    "\n",
    "# To load a dataset\n",
    "wiki = load_dataset(\"wikipedia\", \"20200501.en\", split='train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_system_id = '<FSX_ID>' # FSx file system ID with your training dataset. Example: 'fs-0bYYYYYY'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "SM_DATA_ROOT = '/opt/ml/input/data/train'\n",
    "\n",
    "hyperparameters={\n",
    "    \"train_dir\": '/'.join([SM_DATA_ROOT, 'tfrecords/train/max_seq_len_128_max_predictions_per_seq_20_masked_lm_prob_15']),\n",
    "    \"val_dir\": '/'.join([SM_DATA_ROOT, 'tfrecords/validation/max_seq_len_128_max_predictions_per_seq_20_masked_lm_prob_15']), \n",
    "    \"log_dir\": '/'.join([SM_DATA_ROOT, 'checkpoints/bert/logs']), \n",
    "    \"checkpoint_dir\": '/'.join([SM_DATA_ROOT, 'checkpoints/bert']), \n",
    "    \"load_from\": \"scratch\", \n",
    "    \"model_type\": \"bert\", \n",
    "    \"model_size\": \"large\", \n",
    "    \"per_gpu_batch_size\": 64, \n",
    "    \"max_seq_length\": 128,\n",
    "    \"max_predictions_per_seq\": 20, \n",
    "    \"optimizer\": \"lamb\", \n",
    "    \"learning_rate\": 0.005, \n",
    "    \"end_learning_rate\": 0.0003, \n",
    "    \"hidden_dropout_prob\": 0.1, \n",
    "    \"attention_probs_dropout_prob\": 0.1,\n",
    "    \"gradient_accumulation_steps\": 1,\n",
    "    \"learning_rate_decay_power\": 0.5, \n",
    "    \"warmup_steps\": 2812, \n",
    "    \"total_steps\": 2000, \n",
    "    \"log_frequency\": 10,\n",
    "    \"run_name\" : job_name,\n",
    "    \"squad_frequency\": 0\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = TensorFlow(entry_point='albert/run_pretraining.py',\n",
    "                        role=role,\n",
    "                        image_uri=docker_image,\n",
    "                        source_dir='deep-learning-models/models/nlp',\n",
    "                        framework_version='2.3.1',\n",
    "                        py_version='py3',\n",
    "                        instance_count=instance_count,\n",
    "                        instance_type=instance_type,\n",
    "                        sagemaker_session=sagemaker_session,\n",
    "                        subnets=subnets,\n",
    "                        hyperparameters=hyperparameters,\n",
    "                        security_group_ids=security_group_ids,\n",
    "                        debugger_hook_config=False,\n",
    "                        # Training using SMDataParallel Distributed Training Framework\n",
    "                        distribution={'smdistributed':{\n",
    "                                        'dataparallel':{\n",
    "                                                'enabled': True\n",
    "                                            }\n",
    "                                        }\n",
    "                                      }\n",
    "                      )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configure FSx Input for the SageMaker Training Job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.inputs import FileSystemInput\n",
    "\n",
    "#YOUR_MOUNT_PATH_FOR_TRAINING_DATA # NOTE: '/fsx/' will be the root mount path. Example: '/fsx/albert''''\n",
    "file_system_directory_path='/fsx/' \n",
    "file_system_access_mode='rw'\n",
    "file_system_type='FSxLustre'\n",
    "\n",
    "train_fs = FileSystemInput(file_system_id=file_system_id,\n",
    "                                    file_system_type=file_system_type,\n",
    "                                    directory_path=file_system_directory_path,\n",
    "                                    file_system_access_mode=file_system_access_mode)\n",
    "data_channels = {'train': train_fs}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit SageMaker training job\n",
    "estimator.fit(inputs=data_channels, job_name=job_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_tensorflow_p36",
   "language": "python",
   "name": "conda_tensorflow_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
